from textrank4zh import TextRank4Keyword, TextRank4Sentence
from fastapi import APIRouter, Depends

from app.models import Req, CommonResult, ResponseCode

router = APIRouter(prefix="/textrank")


# 创建依赖项函数
def get_tr4w():
    return TextRank4Keyword()


def get_tr4s():
    return TextRank4Sentence()


def analyze_text(tr, text, lower=True, window=2, source=None):
    # 检查text是否为空
    if not text:
        raise ValueError("text cannot be empty")

    if isinstance(tr, TextRank4Keyword):
        tr.analyze(text=text, lower=lower, window=window)
    elif isinstance(tr, TextRank4Sentence):
        tr.analyze(text=text, lower=lower, source=source if source is not None else 'all_filters')


@router.post("/keywords")
async def text_rank_keyword(text: Req, tr4w: TextRank4Keyword = Depends(get_tr4w)):
    """
    提取关键词

    参数:
    text: Req - 包含要分析的文本的请求对象
    tr4w: TextRank4Keyword - TextRank4Keyword的实例

    返回:
    List[str] - 提取的关键词列表
    """
    try:
        analyze_text(tr4w, text.text)
    except ValueError as e:
        return CommonResult(code=ResponseCode.BAD_REQUEST, message=str(e))

    keywords = [item.word for item in tr4w.get_keywords(5, word_min_len=2)]
    return CommonResult(code=ResponseCode.SUCCESS, data={"keywords": keywords}, message="Success")


@router.post("/keyphrases")
async def text_rank_keyphrases(text: Req, tr4w: TextRank4Keyword = Depends(get_tr4w)):
    """
    提取关键短语

    参数:
    text: Req - 包含要分析的文本的请求对象
    tr4w: TextRank4Keyword - TextRank4Keyword的实例

    返回:
    List[str] - 提取的关键短语列表
    """
    try:
        analyze_text(tr4w, text.text, window=3)
    except ValueError as e:
        return CommonResult(code=ResponseCode.BAD_REQUEST, message=str(e))

    keyphrases = [phrase for phrase in tr4w.get_keyphrases(keywords_num=20, min_occur_num=2)]
    return CommonResult(code=ResponseCode.SUCCESS, data={"keyphrases": keyphrases}, message="Success")


@router.post("/sentences")
async def text_rank_sentences(text: Req, tr4s: TextRank4Sentence = Depends(get_tr4s)):
    """
    提取关键句子

    参数:
    text: Req - 包含要分析的文本的请求对象
    tr4s: TextRank4Sentence - TextRank4Sentence的实例

    返回:
    List[str] - 提取的关键句子列表
    """
    try:
        analyze_text(tr4s, text.text, source='all_filters')
    except ValueError as e:
        return CommonResult(code=ResponseCode.BAD_REQUEST, message=str(e))

    sentences = [item.sentence for item in tr4s.get_key_sentences(num=3)]
    return CommonResult(code=ResponseCode.SUCCESS, data={"sentences": sentences}, message="Success")
